import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from monai import transforms
import json
import os
import clip
import SimpleITK as sitk

import h5py

#### 用于 ktGaussianSampling
def randp(prob, seed, n_tmp, dim):
    """
    模拟 randp 函数，根据概率分布进行采样
    """
    np.random.seed(seed)
    indices = np.random.choice(len(prob), n_tmp, p=prob / np.sum(prob), replace=False)
    return indices

def ktdup(ph, ti, ny, nt):
    """
    模拟 ktdup 函数，处理重复项
    """
    unique_indices = []
    for i in range(len(ph)):
        index = (ph[i], ti[i])
        if index not in unique_indices:
            unique_indices.append(index)
    new_ph = np.array([idx[0] for idx in unique_indices])
    new_ti = np.array([idx[1] for idx in unique_indices])
    return new_ph, new_ti

def zpad(arr, new_shape):
    """
    模拟 zpad 函数，进行零填充
    """
    padded = np.zeros(new_shape, dtype=arr.dtype)
    current_shape = arr.shape
    slices = [slice(0, min(current_shape[i], new_shape[i])) for i in range(len(current_shape))]
    padded[tuple(slices)] = arr
    return padded

### helper
def read_matlab_v73_file(file_path):
    with h5py.File(file_path, 'r') as file:
        key_list = [key for key in file.keys()]        
        return file[key_list[0]][()]


def datafold_read(data_list):
    with open(data_list) as f:
        json_data = json.load(f)
    return json_data


def convert_to_multi_channel(y):
    y_wt, y_tc, y_et = y > 0, ((y == 1) + (y == 3)) > 0, y == 3
    return torch.stack([y_wt, y_tc, y_et], dim=0)
    

class Task1(Dataset):
    def __init__(self, data_list, cls='lax', phase='train', mask_type='task1_u4'):
        super(Task1, self).__init__()
        """
            cls = 'lax' or 'lvot' or 'sax'
        """
        self.cls = cls
        
        if self.cls == 'lax':
            self.size = [204, 448]
        elif self.cls == 'lvot':
            self.size = [168, 416]
        elif self.cls == 'sax':
            self.size = [246, 512]
        else:
            raise f"data type must be lax, lvot, or sax, but get {self.cls}"
        
        self.key = str('full_sample_' + cls)
        
        self.mask_type = mask_type
        
        self.data_list = data_list[phase]

        self.train_transform = transforms.Compose(
            [
                transforms.ToTensord(keys=['k_space']),
            ]
        )
        
        self.val_transform = transforms.Compose(
            [
                transforms.ToTensord(keys=['k_space']),
            ]
        )
        
        self.test_transform = transforms.Compose(
            [
                transforms.ToTensord(keys=['k_space']),
            ]
        )

        self.phase = phase

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, item):
        data = self.get_data(self.data_list[item])
        
        if self.phase == 'train':
            data = self.train_transform(data)
        elif self.phase == 'val':
            data = self.val_transform(data)
        elif self.phase == 'test':
            data = self.test_transform(data)

        return data
    
    def load_mat(self, path):
        mat_data = read_matlab_v73_file(path)
        return mat_data
    
    def padding(self, x, complex=False):
        len_x = len(x.shape)
        data_shape = x.shape[-2:]

        pad_top = (self.size[0] - data_shape[0]) // 2
        pad_bottom = self.size[0] - data_shape[0] - pad_top
        pad_left = (self.size[1] - data_shape[1]) // 2
        pad_right = self.size[1] - data_shape[1] - pad_left

        pad_list = [(0, 0)] * (len_x - 2)
        pad_list = pad_list + [(pad_top, pad_bottom), (pad_left, pad_right)]
        pad_width = (*pad_list, )
        
        if complex:
            return np.pad(x, pad_width, mode='constant', constant_values=0 + 0j)
        else:
            return np.pad(x, pad_width, mode='constant', constant_values=0)        
        
    def get_data(self, file_dic=None):
        indice = file_dic['indice']

        # 获取 K 空间数据，返回四维数组，第一维为时间与相位叠加，c 为线圈， hw 空间分辨率
        full_sample_path = file_dic[self.key]
        
        full_sample = self.load_mat(full_sample_path)
        t, p, c, h, w = full_sample.shape 

        full_sample = np.reshape(full_sample, (t * p, c, h, w))
        
        real_part = full_sample['real']
        imag_part = full_sample['imag']
        
        full_sample_complex = real_part + 1j * imag_part
        full_sample_complex_padded = self.padding(full_sample_complex, complex=True)

        return {
            'k_space': full_sample_complex_padded,
            'indice': indice,
        }
        
def get_mask(pattern='Uniform', R=4, nx=204, ny=448, nt=12, ncalib=16):
    """
        pattern: 'Uniform', 'KGaussion', 
        R: 4, 8, 10, 
    """
    if pattern == 'Uniform':
        # 生成均匀采样掩码
        mask = np.zeros((nx, ny))
        calib_start = (nx - ncalib) // 2
        calib_end = calib_start + ncalib
        mask[calib_start:calib_end, :] = 1  # 校准区域
        # 非校准区域均匀采样
        remaining_indices = np.setdiff1d(np.arange(nx), np.arange(calib_start, calib_end))
        num_samples = int((nx - ncalib) / R)
        sampled_indices = np.random.choice(remaining_indices, num_samples, replace=False)
        mask[sampled_indices, :] = 1
    return torch.from_numpy(mask)       
        
def get_loader(data_list_json,
               batch_size,
               num_works,
               cls='lax',
               phase='train'):

    files = datafold_read(data_list=data_list_json)
    
    datasets = Task1(data_list=files, phase=phase, cls=cls)

    if phase != 'train':
        dataloader = DataLoader(datasets,
                                batch_size=batch_size,
                                num_workers=num_works,
                                pin_memory=True,
                                shuffle=False,
                                drop_last=True)
    else:
        dataloader = DataLoader(datasets,
                                batch_size=batch_size,
                                num_workers=num_works,
                                pin_memory=True,
                                shuffle=True,
                                drop_last=True)
    return dataloader
    
if __name__ == "__main__":    
    # path = "/dev/raid/doctor/wyx/ChallengeData/home2/Raw_data/MICCAIChallenge2024/ChallengeData/MultiCoil/Cine/TrainingSet/Mask_Task1/P016/cine_lvot_mask_Uniform4.mat"
    # mat = read_matlab_v73_file(path)
    # print(mat, mat.shape)
    # assert 1 == 2
    
    data_list = '/dev/raid/doctor/wyx/ChallengeData/home2/Raw_data/MICCAIChallenge2024/ChallengeData/MultiCoil/Cine/TrainingSet/lax_ks.json'
    
    dataloader = get_loader(
        data_list_json=data_list,
        batch_size=4,
        num_works=0,
        cls='lax',
        phase='train'
    )
    
    for i in dataloader:
        k_space = i['k_space']      
        b, n, c, h, w = k_space.size()
        k_space = k_space.view(b*n, c, h, w)
        mask = get_mask()
        
        print("batch finished", k_space.size(), mask.size())
        


    
        

